package com.dragon.aipictures.manager.sharding.test;

import cn.hutool.core.lang.ConsistentHash;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingValue;

import java.util.*;

public class PictureDatabaseShardingAlgorithm implements ComplexKeysShardingAlgorithm<Long> {
    
    // 一致性哈希环（虚拟节点数=1000）
    private static final ConsistentHash<Integer> CONSISTENT_HASH =
        new ConsistentHash<>(1000, Arrays.asList(0,1,2));

    @Override
    public Collection<String> doSharding(Collection<String> availableTargets,
                                         ComplexKeysShardingValue<Long> shardingValue) {
        Map<String, Collection<Long>> columns = shardingValue.getColumnNameAndShardingValuesMap();
        Long spaceId = columns.get("spaceId").iterator().next();

       // 动态分库规则,如果spaceid≠0,根据空间id进行路由
        if (spaceId!=0) {
            int dbIndex = CONSISTENT_HASH.get(spaceId); // 一致性哈希分配
            return Collections.singleton("ai_pictures" + dbIndex);
        }

        // 静态分库规则,如果（spaceid=0），公共空间
        Long pictureId = columns.get("pictureId").iterator().next();
        int dbIndex = (int)(pictureId % 3);
        return Collections.singleton("ai_pictures" + dbIndex);
    }

    @Override
    public Properties getProps() {
        return null;
    }

    @Override
    public void init(Properties properties) {

    }
}